#pragma once

#include "udp_define.hh"
#include <string>
#include <cstdlib>
#include <ctime>

namespace kratos {
namespace network {

#if (defined(WIN32) || defined(_WIN64))
static inline bool recvfrom(SocketType socket, char* buffer, std::size_t& length, SocketAddress& address) {
    WSABUF wsaBuffer = { (ULONG)length, buffer };
    INT addrLength = sizeof(SocketAddress);
    length = 0;
    DWORD flags = 0;
    auto error = ::WSARecvFrom(socket, &wsaBuffer, 1, (LPDWORD)(&length), &flags,
        (struct sockaddr *)&address, &addrLength, nullptr, nullptr);
    if (SOCKET_ERROR == error) {
        switch (WSAGetLastError())
        {
        case WSAEWOULDBLOCK:
        case WSAECONNRESET:
            return true;
        default:
            return false;
        }
    }
    return true;
}
#else
static inline bool recvfrom(SocketType socket, char* buffer, std::size_t& length, SocketAddress& address) {
    socklen_t fromlen = sizeof(address);
    auto bytes = ::recvfrom(socket, buffer, length, 0, (struct sockaddr*)&address, &fromlen);
    if (bytes < 0) {
        length = 0;
        if (errno == EWOULDBLOCK) {
            return true;
        } else {
            return false;
        }
    }
    length = bytes;
    return true;
}
#endif /* defined(WIN32) || defined(_WIN64) */

#if (defined(WIN32) || defined(_WIN64))
static inline bool sendto(SocketType socket, const char* buffer, std::size_t& length, SocketAddress& address) {
    WSABUF wsaBuffer = { (ULONG)length, const_cast<char*>(buffer) };
    length = 0;
    auto error = ::WSASendTo(socket, &wsaBuffer, 1, (LPDWORD)(&length), 0,
        (struct sockaddr *)&address, sizeof(SocketAddress), nullptr, nullptr);
    if (SOCKET_ERROR == error) {
        if (::WSAGetLastError() == WSAEWOULDBLOCK) {
            return true;
        }
        else {
            return false;
        }
    }
    return true;
}
#else
static inline bool sendto(SocketType socket, const char* buffer, std::size_t& length, SocketAddress& address) {
    auto bytes = ::sendto(socket, buffer, length, 0, (struct sockaddr*)&address, sizeof(address));
    if (bytes < 0) {
        length = 0;
        if (errno == EWOULDBLOCK) {
            return true;
        } else {
            return false;
        }
    }
    length = bytes;
    return true;
}
#endif /* defined(WIN32) || defined(_WIN64) */

#if (defined(WIN32) || defined(_WIN64))
static inline void closesocket(SocketType& socket) {
    if (INVALID_SOCKET != socket) {
        ::closesocket(socket);
        socket = INVALID_SOCKET;
    }
}
#else
static inline void closesocket(SocketType& socket) {
    ::close(socket);
}
#endif /* defined(WIN32) || defined(_WIN64) */

static inline bool isSocket(SocketType socket) {
#if (defined(WIN32) || defined(_WIN64))
    return ((INVALID_SOCKET != socket) && (socket > 0));
#else
    return (socket > 0);
#endif /* defined(WIN32) || defined(_WIN64) */
}

static inline SocketType socket() {
    return ::socket(PF_INET, SOCK_DGRAM, 0);
}

static inline void setNonblocking(SocketType socket) {
#if (defined(WIN32) || defined(_WIN64))
    u_long on = 1;
    ::ioctlsocket(socket, FIONBIO, &on);
#else
    int on = 1;
    ::ioctl(socket, FIONBIO, &on);
#endif // #if (defined(WIN32) || defined(_WIN64))
}

static inline bool bind(SocketType socket, const std::string & ip, std::uint16_t port, SocketAddress& address) {
    memset(&address, 0, sizeof(address));
    address.sin_family = AF_INET;
    address.sin_port = htons(port);
#if (defined(WIN32) || defined(_WIN64))
    address.sin_addr.S_un.S_addr = inet_addr(ip.c_str());
#else
    address.sin_addr.s_addr = inet_addr(ip.c_str());
#endif /* defined(WIN32) || defined(_WIN64) */
    auto error = ::bind(socket, (struct sockaddr*)&address, sizeof(address));
    if (error < 0) {
        closesocket(socket);
        return false;
    }
    return true;
}

static inline bool bind(SocketType socket, SocketAddress& address) {
    struct sockaddr_in sin;
    memset(&sin, 0, sizeof(struct sockaddr_in));
    sin.sin_family = AF_INET;
    sin.sin_port = 0;
#if (defined(WIN32) || defined(_WIN64))
    sin.sin_addr.S_un.S_addr = htonl(INADDR_ANY);
#else
    sin.sin_addr.s_addr = htonl(INADDR_ANY);
#endif /* defined(WIN32) || defined(_WIN64) */
    auto error = ::bind(socket, (struct sockaddr*)&sin, sizeof(sin));
    if (error < 0) {
        closesocket(socket);
        return false;
    }
    return true;
}

static inline void udp_os_initialize() {
#if defined(WIN32) || defined(WIN64)
    WSADATA wsa;
    WSAStartup(MAKEWORD(2, 2), &wsa);
#endif /* defined(WIN32) || defined(WIN64) */
}

static inline void udp_os_deinitialize() {
#if defined(WIN32) || defined(WIN64)
    WSACleanup();
#endif /* defined(WIN32) || defined(WIN64) */
}

#if defined(WIN32) || defined(WIN64)
static inline DWORD getLastError() {
    return ::GetLastError();
}
#else
static inline int getLastError() {
    return errno;
}
#endif /* defined(WIN32) || defined(WIN64) */

static inline bool select(SocketType socket, std::time_t million, bool triggerWrite, bool& isRead, bool& isWrite) {
    fd_set readSet, writeSet;
    struct timeval tv = { (long)million / 1000, ((long)million % 1000) * 1000 };
    FD_ZERO(&readSet);
    FD_ZERO(&writeSet);
    if (triggerWrite) {
        FD_SET(socket, &writeSet);
    }
    FD_SET(socket, &readSet);
    auto count = ::select((int)socket + 1, &readSet, &writeSet, nullptr, &tv);
    if (count <= 0) {
        return false;
    }
    if (FD_ISSET(socket, &writeSet)) {
        isWrite = true;
    }
    if (FD_ISSET(socket, &readSet)) {
        isRead = true;
    }
    return isWrite || isRead;
}

/// UDP address
class UdpAddress {
    std::uint64_t id_; // 64bit ID
    SocketAddress address_; // address
public:
    // ctor
    UdpAddress()
        : id_(0) {
    }
    // ctor
    // @param ip IP
    // @param port port
    UdpAddress(const std::string& ip, std::uint16_t port)
        : id_(0) {
        setAddress(ip, port);
    }
    // copy ctor
    // @param address other instance
    UdpAddress(SocketAddress& address) {
        address_ = address;
#if (defined(WIN32) || defined(_WIN64))
        uint64_t dword = address.sin_addr.S_un.S_addr;
#else
        uint64_t dword = address.sin_addr.s_addr;
#endif /* defined(WIN32) || defined(_WIN64) */
        id_ = dword << 32;
        id_ += (std::uint64_t)address.sin_port;
    }
    // dtor
    ~UdpAddress() {
    }
    // assigns address
    // @param address other instance
    void setAddress(const SocketAddress& address) {
        address_ = address;
#if (defined(WIN32) || defined(_WIN64))
        uint64_t dword = address.sin_addr.S_un.S_addr;
#else
        uint64_t dword = address.sin_addr.s_addr;
#endif /* defined(WIN32) || defined(_WIN64) */
        id_ = dword << 32;
        id_ += (std::uint64_t)address.sin_port;
    }
    // assign address
    // @param ip IP
    // @param port port
    void setAddress(const std::string& ip, std::uint16_t port) {
        memset(&address_, 0, sizeof(SocketAddress));
        address_.sin_family = AF_INET;
        address_.sin_port = htons(port);
#if (defined(WIN32) || defined(_WIN64))
        address_.sin_addr.S_un.S_addr = inet_addr(ip.c_str());
#else
        address_.sin_addr.s_addr = inet_addr(ip.c_str());
#endif /* defined(WIN32) || defined(_WIN64) */
#if (defined(WIN32) || defined(_WIN64))
        uint64_t dword = inet_addr(ip.c_str());
#else
        uint64_t dword = inet_addr(ip.c_str());
#endif /* defined(WIN32) || defined(_WIN64) */
        id_ = (std::uint64_t)dword << 32;
        id_ += (std::uint64_t)port;
    }
    // returns address
    SocketAddress& getAddress() {
        return address_;
    }
    // returns 64bit ID
    std::uint64_t getID() {
        return id_;
    }
};

}
}
